// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <turbo/log/logging.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/memory/memory.h>
#include <pollux/dwio/common/data_buffer.h>

namespace kumo {
namespace pollux {
namespace dwio {
namespace common {
using namespace kumo::pollux::memory;
using namespace testing;
using MemoryPool = kumo::pollux::memory::MemoryPool;

class DataBufferTest : public testing::Test {
 protected:
  static void SetUpTestCase() {
    MemoryManager::testingSetInstance({});
  }

  const std::shared_ptr<MemoryPool> pool_ = memoryManager()->addLeafPool();
};

TEST_F(DataBufferTest, ZeroOut) {
  const uint8_t VALUE = 13;
  DataBuffer<uint8_t> buffer(*pool_, 16);
  for (auto i = 0; i < buffer.size(); i++) {
    auto data = buffer.data();
    ASSERT_EQ(data[i], 0);
    data[i] = VALUE;
  }

  buffer.resize(8);
  for (auto i = 0; i < buffer.size(); i++) {
    auto data = buffer.data();
    ASSERT_EQ(data[i], VALUE);
  }

  auto currentSize = buffer.size();
  ASSERT_EQ(currentSize, 8);
  buffer.resize(32);
  for (auto i = 0; i < buffer.size(); i++) {
    auto data = buffer.data();
    if (i < currentSize) {
      ASSERT_EQ(data[i], VALUE);
    } else {
      ASSERT_EQ(data[i], 0);
    }
  }
}

TEST_F(DataBufferTest, At) {
  DataBuffer<uint8_t> buffer{*pool_};
  for (auto i = 0; i != 15; ++i) {
    buffer.append(i);
  }
  ASSERT_EQ(15, buffer.size());

  for (auto i = 0; i != 15; ++i) {
    EXPECT_EQ(i, buffer.at(i));
  }

  buffer.resize(8);
  for (auto i = 0; i != 8; ++i) {
    EXPECT_EQ(i, buffer.at(i));
  }
  for (auto i = 8; i != 42; ++i) {
    POLLUX_ASSERT_THROW(buffer.at(i), "Accessing index out of range");
  }
}

TEST_F(DataBufferTest, Reset) {
  DataBuffer<uint8_t> buffer{*pool_};
  buffer.reserve(16);
  for (auto i = 0; i != 15; ++i) {
    buffer.append(i);
  }
  ASSERT_EQ(15, buffer.size());
  ASSERT_EQ(16, buffer.capacity());
  {
    buffer.clear();
    EXPECT_EQ(0, buffer.size());
    EXPECT_EQ(0, buffer.capacity());

    buffer.reserve(12);
    EXPECT_EQ(0, buffer.size());
    EXPECT_EQ(12, buffer.capacity());
    for (auto i = 0; i != 11; ++i) {
      buffer.append(i);
    }
    EXPECT_EQ(11, buffer.size());
    EXPECT_EQ(12, buffer.capacity());
  }

  {
    buffer.clear();
    EXPECT_EQ(0, buffer.size());
    EXPECT_EQ(0, buffer.capacity());

    buffer.reserve(16);
    EXPECT_EQ(0, buffer.size());
    EXPECT_EQ(16, buffer.capacity());
    for (auto i = 0; i != 15; ++i) {
      buffer.append(i);
    }
    EXPECT_EQ(15, buffer.size());
    EXPECT_EQ(16, buffer.capacity());
  }

  {
    buffer.clear();
    EXPECT_EQ(0, buffer.size());
    EXPECT_EQ(0, buffer.capacity());

    buffer.reserve(32);
    EXPECT_EQ(0, buffer.size());
    EXPECT_EQ(32, buffer.capacity());
    for (auto i = 0; i != 31; ++i) {
      buffer.append(i);
    }
    EXPECT_EQ(31, buffer.size());
    EXPECT_EQ(32, buffer.capacity());
  }
}

TEST_F(DataBufferTest, Wrap) {
  auto size = 26;
  auto buffer = pollux::AlignedBuffer::allocate<char>(size, pool_.get());
  auto raw = buffer->asMutable<char>();
  for (size_t i = 0; i < size; ++i) {
    raw[i] = 'a' + i;
  }
  auto dataBuffer = DataBuffer<char>::wrap(buffer);
  buffer = nullptr;
  ASSERT_EQ(size, dataBuffer->size());
  ASSERT_EQ(size, dataBuffer->capacity());
  for (size_t i = 0; i < size; ++i) {
    ASSERT_EQ((*dataBuffer)[i], 'a' + i);
  }
}

TEST_F(DataBufferTest, Move) {
  {
    DataBuffer<uint8_t> buffer{*pool_};
    buffer.reserve(16);
    for (auto i = 0; i != 15; ++i) {
      buffer.append(i);
    }
    ASSERT_EQ(15, buffer.size());
    ASSERT_EQ(16, buffer.capacity());
    const auto usedBytes = pool_->usedBytes();

    // Expect no double freeing from memory pool.
    DataBuffer<uint8_t> newBuffer{std::move(buffer)};
    ASSERT_EQ(15, newBuffer.size());
    ASSERT_EQ(16, newBuffer.capacity());
    ASSERT_EQ(usedBytes, pool_->usedBytes());
  }
  ASSERT_EQ(0, pool_->usedBytes());
}
} // namespace common
} // namespace dwio
} // namespace pollux
} // namespace kumo
